import time
from wxauto import WeChat
from ollama import OllamaAPI  # 导入 OllamaAPI 类

class OllamaAPIWithPersistence(OllamaAPI):
    def __init__(self, config_file='config.json', history_file='user_messages.json'):
        super().__init__(config_file, history_file)

wx = WeChat()
ollama = OllamaAPIWithPersistence()  # 使用带有持久化功能的 OllamaAPI 实例

while True:
    msgs = wx.GetNextNewMessage()  # 获取新消息

    if msgs:
        for sender, messages in msgs.items():
            for msg in messages:
                sender_name = msg[0]  # 发送者
                content = msg[1]  # 消息内容
                
                print(f"收到来自 [{sender_name}] 的消息：{content}")

                # 如果是 "SYS" 消息，则不回复
                if sender_name == "SYS" or sender_name == "Self":
                    continue

                # 获取该用户的自动回复，传入 sender_name 来区分上下文
                response = ollama.get_reply(sender_name, content)

                # 如果 Ollama 有回复，发送回复消息
                if response:
                    wx.SendMsg(response, sender_name)
                    print(f"已回复 [{sender_name}]：{response}")
                else:
                    wx.SendMsg("对不起，无法获取回复", sender_name)
                    print(f"已回复 [{sender_name}]：对不起，无法获取回复")

    else:
        print("当前没有新消息")
    
    time.sleep(3)  # 休眠 3 秒